package com.bittech.facerecognition.client;

import com.baidu.aip.face.AipFace;
import com.baidu.aip.face.MatchRequest;
import com.bittech.facerecognition.config.FaceModel;
import org.json.JSONObject;

import java.util.Arrays;
import java.util.Base64;
import java.util.HashMap;

/**
 * Author: secondriver
 * Created: 2019/5/28
 */
public class FaceAiClient {
    
    private static final String APP_ID = "9613825";
    private static final String API_KEY = "lYy9MBOjR9YwZa8mf9R3aq3q";
    private static final String SECRET_KEY = "jG3FY5CZOUdkAXpqhhUxHuNw6V4UKE1b";
    private static final String IMAGE_TYPE = "BASE64";
    
    private final AipFace aipFace;
    
    public FaceAiClient() {
        aipFace = new AipFace(APP_ID, API_KEY, SECRET_KEY);
        aipFace.setConnectionTimeoutInMillis(2000);
        aipFace.setSocketTimeoutInMillis(60000);
    }
    
    public FaceModel detect(byte[] img) {
        HashMap<String, String> options = new HashMap<>();
        options.put("face_field", "age,beauty,expression,faceshape,gender,glasses,landmark,race,quality,facetype");
        options.put("max_face_num", "1");
        JSONObject jsonObject = aipFace.detect(bytesToBase64(img), IMAGE_TYPE, options);
        System.out.println(jsonObject.toString());
        if (jsonObject.optString("error_msg", "").equalsIgnoreCase("SUCCESS")) {
            JSONObject resultJsonObject = jsonObject.optJSONObject("result");
            if (resultJsonObject.optInt("face_num") == 1) {
                JSONObject faceObject = (JSONObject) resultJsonObject.optJSONArray("face_list").get(0);
                return parseAttribute(faceObject);
            }
        }
        return null;
    }
    
    
    private FaceModel parseAttribute(JSONObject faceObject) {
        FaceModel faceModel = new FaceModel();
        //表情
        String expression = faceObject.getJSONObject("expression").getString("type");
        switch (expression) {
            case "smile":
                faceModel.setExpression("微笑");
                break;
            case "laugh":
                faceModel.setExpression("大笑");
                break;
            case "none":
            default:
                faceModel.setExpression("无");
        }
        
        //脸型
        String faceShape = faceObject.getJSONObject("face_shape").getString("type");
        switch (faceShape) {
            case "square":
                faceModel.setFaceShape("正方形");
                break;
            case "triangle":
                faceModel.setFaceShape("三角形");
                break;
            case "oval":
                faceModel.setFaceShape("椭圆");
                break;
            case "heart":
                faceModel.setFaceShape("心形");
                break;
            case "round":
                faceModel.setFaceShape("圆形");
                break;
            default:
                faceModel.setFaceShape("未知");
            
        }
        
        //性别
        String gender = faceObject.getJSONObject("gender").getString("type");
        switch (gender) {
            case "male":
                faceModel.setGender("男性");
                break;
            case "female":
                faceModel.setGender("女性");
                break;
            default:
                faceModel.setGender("未知");
        }
        
        //人种
        String race = faceObject.getJSONObject("race").getString("type");
        switch (race) {
            case "yellow":
                faceModel.setRace("黄种人");
                break;
            case "white":
                faceModel.setRace("白种人");
                break;
            case "black":
                faceModel.setRace("黑种人");
                break;
            case "arabs":
                faceModel.setRace("阿拉伯人");
                break;
            default:
                faceModel.setRace("未知");
        }
        
        //眼镜
        String glasses = faceObject.getJSONObject("glasses").getString("type");
        switch (glasses) {
            case "common":
                faceModel.setGlasses("普通眼镜");
                break;
            case "sun":
                faceModel.setGlasses("墨镜");
                break;
            case "none":
                faceModel.setGlasses("无眼镜");
                break;
            default:
                faceModel.setGlasses("未知");
        }
        
        //年龄
        int age = faceObject.optInt("age");
        faceModel.setAge(age);
        
        //颜值
        double beauty = faceObject.optDouble("beauty");
        faceModel.setBeauty(beauty);
        
        return faceModel;
    }
    
    public Double match(byte[] img1, byte[] img2) {
        MatchRequest request1 = new MatchRequest(bytesToBase64(img1), IMAGE_TYPE);
        MatchRequest request2 = new MatchRequest(bytesToBase64(img2), IMAGE_TYPE);
        JSONObject jsonObject = aipFace.match(Arrays.asList(request1, request2));
        System.out.println(jsonObject.toString());
        if (jsonObject.optString("error_msg", "").equalsIgnoreCase("SUCCESS")) {
            return jsonObject.optJSONObject("result").optDouble("score", 0D);
        }
        return 0D;
    }
    
    private String bytesToBase64(byte[] data) {
        if (data == null) {
            data = new byte[]{};
        }
        return Base64.getEncoder().encodeToString(data);
    }
    
}
